package org.example.develop;

import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.api.java.UDF2;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.functions;

import java.util.TreeSet;
import java.util.List;
import java.util.stream.Collectors;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/***
 * 自定义UDF,实现对模型一:卡合 模型二:卡分中健康评估涉及到的轴信息补充上来,从而进行健康评估.
 */
public class NumberExtractor {
    private SparkSession spark;

    public NumberExtractor(SparkSession spark) {
        this.spark = spark;
        registerUDF();
    }

    private void registerUDF() {
        spark.udf().register("extract_and_convert_numbers", new ExtractAndConvertNumbersUDF(), DataTypes.StringType);
    }

    public Dataset<Row> applyUDF(Dataset<Row> dataset) {
        return dataset.withColumn("extracted_numbers",
                functions.callUDF("extract_and_convert_numbers", functions.col("phm_gzmc"), functions.col("bjwz").cast(DataTypes.IntegerType)));
    }


    public static class ExtractAndConvertNumbersUDF implements UDF2<String, Integer, String> {
        @Override
        public String call(String phm_gzmc, Integer bjwz) throws Exception {
            // 将phm_gzmc使用逗号分隔，得到一个字符串数组
            String[] items = phm_gzmc.split(",");
            // 创建一个TreeSet，用于存放找到的数字（TreeSet会自动对插入的数字进行排序且去重）
            TreeSet<Integer> numbers = new TreeSet<>();
            // 使用正则表达式匹配一、二、三中的任意一个
            // Pattern.compile("(一|二|三)")编译一个正则表达式模式，该模式可以匹配一、二、三中的任意一个字符
            Pattern pattern = Pattern.compile("(一|二|三)");
            // 遍历items数组中的每个字符串
            for (String item : items) {
                // 使用正则表达式模式在当前字符串中查找匹配项
                Matcher matcher = pattern.matcher(item);
                // 如果找到了匹配的字符（一、二、三中的一个）
                if (matcher.find()) {
                    // 获取匹配到的字符（一、二、三中的一个）
                    // matcher.group(1)表示获取正则表达式中第一个括号内匹配到的内容
                    String matched = matcher.group(1);
                    int number;
                    // 根据匹配到的字符（一、二、三）和bjwz的值，计算对应的数字
                    switch (matched) {
                        case "一":
                            number = bjwz == 1 ? 1 : 6;
                            break;
                        case "二":
                            number = bjwz == 1 ? 2 : 5;
                            break;
                        case "三":
                            number = bjwz == 1 ? 3 : 4;
                            break;
                        default:
                            continue;
                    }
                    numbers.add(number);
                }
            }

            return numbers.stream().map(String::valueOf).collect(Collectors.joining(","));
        }
    }
}
